#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
选股系统 (SQLAlchemy版本)
使用SQLAlchemy进行数据库操作
"""

import pandas as pd
import numpy as np
from datetime import datetime
from typing import List, Dict
from loguru import logger

from Strategy import StockStrategies
from model import DatabaseManager, HSnapshot, Block, BlockInfo, SelectResult, LabelInfo

class StockSelector:
    """使用SQLAlchemy的选股器"""
    
    def __init__(self):
        """初始化选股器"""
        # 初始化SQLAlchemy数据库管理器
        self.db_manager = DatabaseManager()
        self.db_manager.init_databases()
        
        # 设置模型的数据库管理器
        HSnapshot.set_db_manager(self.db_manager)
        Block.set_db_manager(self.db_manager)
        BlockInfo.set_db_manager(self.db_manager)
        SelectResult.set_db_manager(self.db_manager)
        LabelInfo.set_db_manager(self.db_manager)
        
        # 策略列表（统一从StrategyConfig获取）
        from StrategyConfig import get_strategy_list
        self.strategy_list = get_strategy_list()
    
    def load_data(self, db_path: str = "stock_data.db") -> pd.DataFrame:
        """加载股票数据"""
        logger.info("正在加载股票数据...")
        
        # 加载股票数据
        stock_data = HSnapshot.load_from_db(db_path=db_path)
        
        if stock_data.empty:
            logger.warning("数据库中没有找到数据")
            logger.error("无法加载股票数据，请检查数据源")
            return pd.DataFrame()
        
        # 数据验证
        if not stock_data.empty:
            required_columns = ['security_id', 'trade_date', 'close', 'volume']
            missing_columns = [col for col in required_columns if col not in stock_data.columns]
            if missing_columns:
                logger.error(f"数据缺少必要列: {missing_columns}")
                return pd.DataFrame()
        
        # 数据预处理
        stock_data = HSnapshot.preprocess_data(stock_data)
        
        # 数据过滤优化
        if not stock_data.empty:
            stock_data = stock_data[
                (stock_data['close'] > 0) & 
                (stock_data['volume'] > 0) &
                (stock_data['close'] < 100000)
            ]
            logger.info(f"数据过滤后，共 {len(stock_data)} 条记录")
        
        logger.info(f"数据加载完成，共 {len(stock_data)} 条记录")
        return stock_data
    
    def execute_strategies(self, stock_data: pd.DataFrame, execution_date: str = None) -> pd.DataFrame:
        """执行选股策略"""
        if execution_date is None:
            execution_date = datetime.now().strftime('%Y%m%d')
        
        logger.info(f"开始执行选股策略，执行日期: {execution_date}")
        
        # 加载板块数据
        block_data = Block.load_from_db("stock_data.db")
        blockinfo_data = BlockInfo.load_from_db("stock_data.db")
        
        # 执行策略
        all_results = []
        total_strategies = len(self.strategy_list)
        successful_strategies = 0
        
        logger.info(f"将执行 {total_strategies} 个选股策略")
        
        for i, (strategy_name, chinese_name) in enumerate(self.strategy_list, 1):
            logger.info(f"正在执行策略 [{i}/{total_strategies}]: {chinese_name}")
            
            try:
                # 执行策略
                selected_stocks = StockStrategies.apply_strategy_to_all(stock_data, strategy_name)
                
                if not selected_stocks.empty:
                    # 生成标准化输出
                    # 临时处理：直接使用选股结果，不进行标准化
                    standardized_output = selected_stocks
                    
                    if not standardized_output.empty:
                        all_results.append(standardized_output)
                        successful_strategies += 1
                        logger.info(f"策略 {chinese_name} 选出 {len(standardized_output)} 只股票")
                    else:
                        logger.warning(f"策略 {chinese_name} 标准化输出为空")
                else:
                    logger.info(f"策略 {chinese_name} 没有选出股票")
                    
            except Exception as e:
                logger.error(f"执行策略 {chinese_name} 时出错: {str(e)}")
                continue
        
        logger.info(f"策略执行完成: {successful_strategies}/{total_strategies} 个策略成功选出股票")
        
        # 合并结果
        if all_results:
            combined_results = pd.concat(all_results, ignore_index=True)
            
            # 按执行交易日和证券代码排序
            combined_results = combined_results.sort_values(['执行交易日', '证券代码'])
            
            # 去重处理
            combined_results = combined_results.drop_duplicates(subset=['执行交易日', '证券代码'], keep='first')
            logger.info(f"去重后，共选出 {len(combined_results)} 只股票")
            
            return combined_results
        else:
            logger.warning("没有策略选出股票")
            return pd.DataFrame()
    
    def save_results(self, results: pd.DataFrame, execution_date: str = None) -> bool:
        """保存选股结果到数据库"""
        if results.empty:
            logger.warning("没有结果需要保存")
            return False
        
        if execution_date is None:
            execution_date = datetime.now().strftime('%Y%m%d')
        
        logger.info(f"正在保存选股结果到数据库，执行日期: {execution_date}")
        
        try:
            # 准备数据
            results_data = []
            for _, row in results.iterrows():
                results_data.append({
                    'execution_date': row['执行交易日'],
                    'strategy_type': row['选股类型'],
                    'security_id': row['证券代码'],
                    'security_name': row.get('证券名称', ''),
                    'latest_price': row.get('最新价格', 0.0),
                    'change_percent': row.get('涨跌幅(%)', 0.0),
                    'industry': row.get('所属行业', ''),
                    'limit_up_20d': row.get('近20天涨停天数', 0),
                    'limit_up_60d': row.get('近60天涨停天数', 0),
                    'last_limit_up_date': row.get('最近一次涨停日期', ''),
                    'net_inflow_3d': row.get('最近3日净流入额(万元)', 0.0),
                    'big_net_inflow_3d': row.get('最近3日主力净流入额(万元)', 0.0)
                })
            
            # 保存到数据库
            success = SelectResult.save_batch(results_data)
            if success:
                logger.info(f"成功保存 {len(results_data)} 条选股结果到数据库")
                # 同步写入标签信息（将策略中文作为标签，日期为执行日）
                try:
                    label_rows = []
                    # 兼容不同字段命名，优先中文列
                    date_col = '执行交易日' if '执行交易日' in results.columns else 'execution_date'
                    sid_col = '证券代码' if '证券代码' in results.columns else 'security_id'
                    label_col = '选股类型' if '选股类型' in results.columns else 'strategy_type'
                    for _, row in results.iterrows():
                        label_rows.append({
                            'date': row.get(date_col, execution_date),
                            'security_id': str(row.get(sid_col, '')),
                            'label': str(row.get(label_col, '')),
                            'label_type': 'strategy',
                            'confidence': 100
                        })
                    if label_rows:
                        df_labels = pd.DataFrame(label_rows)
                        saved_cnt = LabelInfo.import_from_dataframe(df_labels)
                        logger.info(f"已写入/更新 {saved_cnt} 条标签记录到 LabelInfo")
                except Exception as e:
                    logger.error(f"写入标签信息失败: {str(e)}")
                return True
            else:
                logger.error("保存选股结果到数据库失败")
                return False
                
        except Exception as e:
            logger.error(f"保存到数据库时出错: {str(e)}")
            return False
    
    def get_available_dates(self) -> List[str]:
        """获取可用的执行日期"""
        # 注意：这个方法在新的模型中可能不存在，需要临时处理
        return []
    
    def get_available_strategies(self) -> List[str]:
        """获取可用的策略类型"""
        # 注意：这个方法在新的模型中可能不存在，需要临时处理
        return []
    
    def get_selection_results(self, execution_date: str = None, strategy_types: List[str] = None, limit: int = 1000):
        """获取选股结果"""
        return SelectResult.get_results(
            execution_date=execution_date,
            strategy_types=strategy_types,
            limit=limit
        )
    
    def clear_cache(self, execution_date: str, strategy_types: List[str] = None):
        """清空缓存"""
        session = self.db_manager.get_stock_selection_session()
        try:
            from model import SelectResult, ExecutionCache
            
            if strategy_types:
                # 删除指定策略的结果
                deleted_results = session.query(SelectResult).filter(
                    SelectResult.execution_date == execution_date,
                    SelectResult.strategy_type.in_(strategy_types)
                ).delete()
                
                # 删除执行缓存
                deleted_cache = session.query(ExecutionCache).filter(
                    ExecutionCache.execution_date == execution_date,
                    ExecutionCache.strategy_type.in_(strategy_types)
                ).delete()
            else:
                # 删除指定日期的所有结果
                deleted_results = session.query(SelectResult).filter(
                    SelectResult.execution_date == execution_date
                ).delete()
                
                # 删除执行缓存
                deleted_cache = session.query(ExecutionCache).filter(
                    ExecutionCache.execution_date == execution_date
                ).delete()
            
            session.commit()
            logger.info(f"清空缓存完成: 删除了 {deleted_results} 条选股结果和 {deleted_cache} 条缓存记录")
            return True
            
        except Exception as e:
            logger.error(f"清空缓存失败: {str(e)}")
            session.rollback()
            return False
        finally:
            session.close()

def main():
    """主函数"""
    import time
    start_time = time.time()
    
    logger.info("开始执行选股程序 (SQLAlchemy版本)")
    
    # 创建选股器
    selector = StockSelector()
    
    # 加载数据
    stock_data = selector.load_data()
    if stock_data.empty:
        logger.error("无法加载数据，程序退出")
        return
    
    # 执行策略
    results = selector.execute_strategies(stock_data)
    if results.empty:
        logger.warning("没有选出股票，程序结束")
        return
    
    # 保存结果
    success = selector.save_results(results)
    if success:
        logger.info("选股结果已成功保存到数据库")
    else:
        logger.error("保存选股结果失败")
    
    # 显示统计信息
    strategy_stats = results['选股类型'].value_counts()
    logger.info("各策略选股统计:")
    for strategy, count in strategy_stats.items():
        logger.info(f"  {strategy}: {count} 只")
    
    # 显示总执行时间
    total_time = time.time() - start_time
    logger.info(f"程序总执行时间: {total_time:.2f} 秒")

if __name__ == "__main__":
    logger.add(
        "stock_select_{time:YYYYMMDD}.log",
        rotation="1 day",
        retention="7 days",
        encoding="utf-8",
        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>"
    )
    main() 